import torch
import torch.nn as nn
from bayesian_base import BayesianBase
import torch.nn.functional as F


class BayesianLinear(BayesianBase):
    """贝叶斯线性层
    Y = x * W^T + b
    where W~N(w_mean_prior, w_std_prior)
          b~N(bias_mean_prior, bias_std_prior)
    y~N(Y, sigma^2)

    Args:
        inputdim ( int ):  输入维度
        outputdim ( int ): 输出维度
        if_bias (bool): 是否包含偏置
    """
    def __init__(self,
                 inputdim,
                 outputdim,
                 if_bias=True):
        super().__init__()
        self.inputdim = inputdim
        self.outputdim = outputdim
        self.if_bias = if_bias
        # initial parameter
        self._init_parameter()
        self.w = None
        # 设置权重w和b(如果包含偏置)的采样器以及先验
        self.normal("w", self.w_mean, self.w_std)
        self.normal("w_prior", mean=0, std=1)
        if self.if_bias:
            self.bias = None
            self.normal("bias", self.bias_mean, self.w_std)
            self.normal("bias_prior", mean=0, std=1)
        #设置y的方差
        self.register_buffer("y_sigma", torch.FloatTensor([1.0]))

    def _init_parameter(self):
        """参数初始化
        """
        self.w_mean = nn.Parameter(
            torch.randn(self.outputdim, self.inputdim).normal_(0, 0.01)
        )
        self.w_std = nn.Parameter(
            torch.randn(self.outputdim, self.inputdim).normal_(0, 0.01)
        )
        if self.if_bias:
            self.bias_mean = nn.Parameter(torch.randn(self.outputdim))
            self.bias_std = nn.Parameter(torch.randn(self.outputdim))

    def forward(self, x):
        self.w = self.node["w"].get_samples
        self.node["w_prior"].set_observation(self.w)
        if self.if_bias:
            self.bias = self.node["bias"].get_samples
            self.node["bias_prior"].set_observation(self.bias)
            y_mean = F.linear(x, self.w, self.bias)
        else:
            y_mean = F.linear(x, self.w)
        self.normal("y_mean", y_mean, self.y_sigma)
        return self.node["y_mean"]
